Codecov Report

❌ Patch coverage is 90.33531% with 49 lines in your changes missing coverage. Please review.

... and 1 file with indirect coverage changes

I can likely pick up some of the linalg ones if you like

I'll keep my progress committed and pushed, feel free to push if you have something. If not I'll just gradually keep adding some whenever I'm waiting for other tests, so also shouldn't be a huge issue.

I definitely agree that it would be nicer to put this in a better form, but would you be okay with leaving that for a follow-up PR?
I tried already separating out some of the functions so it should become easier to migrate this in the future.

What is preventing me from actually pulling this through though is that I now am also altering the primal computation in some places, specifically for constructions involving alpha.
The idea is that for a computation f!(out, args..., alpha, beta) = beta * out + alpha * f(args...) the pullback with respect to alpha is simply derived from f(args...) alone, so I can change the primal computation to f!_mod(out, args..., alpha, beta) = add!(out, f(args...), alpha, beta) and store the intermediate result. I.e. at the cost of adding an additional allocation and an in-place add!, I remove having to compute f in the reverse pass, but only when dalpha is required. (See e.g. the rule for mul!).

Without actually having the Enzyme code next to it, it's a bit hard to already come up with the correct abstractions to make sure this works for both engines, and I want to avoid having to do that work twice.

Additionally, it would be nice to immediately overload the TensorOperations functions but these haven't been released yet (and additionally I would like to play a similar trick there, but haven't gotten around to that yet)

I definitely agree that it would be nicer to put this in a better form, but would you be okay with leaving that for a follow-up PR?

Yeah that sounds fine, just separating things into discrete functions is great already

Small update here:

• This requires another MatrixAlgebraKit release to satisfy the Mooncake 0.5 compat
• I'm adding a custom tangent type here because it turns out my test tolerances were a bit stupid. The finite differences tests that Mooncake is performing are giving wrong answers for non-abelian symmetries in the same way that the ChainRules ones required overloading FiniteDifferences.to_vec, since the inner product would just be the one on the data, rather than the one of the actual tensors.

This last part still really confuses me (I remember it also did with the ChainRules), but I'm just going to assume we figured this out correctly last time and copy that here.

Your PR no longer requires formatting changes. Thank you for your contribution!

Needs QuantumKitHub/MatrixAlgebraKit.jl#165

I think this is ready for a final round of review now, everything seems to be working. @Jutho, thanks for the detailed comments in the previous round, there are definitely a number of things I was missing and would have had a hard time finding myself.